{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change Detection (Segmentation) using TAO Visual ChangeNet-Segmentation Network\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a Python based AI toolkit for taking purpose-built AI models and customizing them with your own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## What is a Visual ChangeNet Network?\n",
    "\n",
    "Visual ChangeNet is a state of the art transformer-based Change Detection model. Visual ChangeNet is based on the Siamese Network, which is a class of neural network architectures containing two or more identical subnetworks. The training algorithm works by updating the parameters across all the sub-networks in tandem. In TAO, Visual ChangeNet supports two images as input where the end goal is to classify or segment the change between the \"golden or reference\" image and the \"test\" image. TAO supports the [FAN](https://arxiv.org/abs/2204.12451) backbone network for both Visual ChangeNet architectures. For more details about training FAN backbones, see the Classification PyTorch notebook.\n",
    "\n",
    "In TAO, two different types of Change Detection networks are supported to detect the difference between the two input images: \n",
    "\n",
    "* Visual ChangeNet-Segmentation for segmentation \n",
    "* Visual ChangeNet-Classification for classification  \n",
    "\n",
    "In this notebook, the Visual ChangeNet-Segmentation model is leveraged to demonstrate change detection by segmenting the changed pixels between the two input images. \n",
    "\n",
    "### Sample Inference from the Visual ChangeNet-Segmentation Model\n",
    "\n",
    "The following is a sample inference output for the Visual ChangeNet-Segmentation model on the MVTec Bottle class. The model takes in two input images (test and golden reference) and outputs a segmentation mask showing the defective region on an object in the MVTec Anomaly Dataset.\n",
    "\n",
    "\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/purpose_built_models/visual_changenet/mvtech_example.png?raw=true\" width=\"1080\">\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you learn how to leverage TAO to:\n",
    "\n",
    "* Train a Visual ChangeNet-Segmentation Model for defect detection on objects from the MVTec Anomaly Dataset.\n",
    "* Evaluate the trained model.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to a .onnx file (encrypted ONNX model) for deployment to DeepStream, TensorRT, or Triton.\n",
    "\n",
    "At the end of this notebook, you generate a trained and optimized `visual_changenet` model, \n",
    "which you can deploy with this [end-to-end sample](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps) with Triton.\n",
    "\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of the Siamese Network using the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset](#head-2)\n",
    "3. [Download the pretrained model from NGC](#head-3)\n",
    "4. [Provide training specification](#head-4)\n",
    "5. [Run TAO training](#head-5)\n",
    "6. [Evaluate trained models](#head-6)\n",
    "7. [Inferences](#head-7)\n",
    "8. [Deploy](#head-8)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The TAO launcher uses Docker containers and **for our data and results directory to be visible to Docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options such as, the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. It maps directories in which the data, specs, results, and cache information is saved. You must configure it for your specific case so that these directories are correctly visible to the Docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=FIXME\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"changenet\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"changenet\", \"results\")\n",
    "os.environ[\"HOST_MODEL_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"changenet\", \"model\")\n",
    "\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR\n",
    "! mkdir -p $HOST_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO Docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_MODEL_DIR\"],\n",
    "           \"destination\": \"/model\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a Python package distributed as a Python Wheel listed in PyPI. Install the launcher by executing the following cell.\n",
    "\n",
    "TAO Toolkit recommends that you run the TAO launcher in a virtual env with Python 3.7.0. Follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a Python virtual env using the `virtualenv` and `virtualenvwrapper` packages. After you setup the virtualenvwrapper, set the version of Python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. To run it:\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 7 and <= 8\n",
    "\n",
    "Perform this step first and then launch the Notebook from the virtual environment. In addition to installing TAO Python package, you must meet the following software requirements:\n",
    "* python >=3.7.0 < 3.8.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 525+\n",
    "\n",
    "After installing the pre-requisites, log in to the Docker registry nvcr.io using the following command:\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "Enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "The TAO Toolkit recommends that you run the TAO launcher in a virtual env with Python >=3.7.0. Follow these [instructions](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a Python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare a Class From The MVTech Anomaly dataset <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "TAO Visual ChangeNet-Segmentation uses a custom dataset format. The sections below walk through this format, specifically the dataset structure and the files required."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual ChangeNet-Segmentation data format\n",
    "\n",
    "Visual ChangeNet-Segmentation expects directories of images and mask files in the dataset root directory. The image directories consist of\n",
    "the golden image directory (pre-change images) and the test image directory (post-change image) to compare against the change mask images with pixel level change masks.  \n",
    "\n",
    "\n",
    "     |--dataset_root:\n",
    "          |--A\n",
    "               |--image1.jpg\n",
    "               |--image2.jpg\n",
    "          |--B\n",
    "               |--image1.jpg\n",
    "               |--image2.jpg\n",
    "          |--label\n",
    "               |--image1.jpg\n",
    "               |--image2.jpg      \n",
    "          |--list\n",
    "               |-- train.txt\n",
    "               |-- val.txt\n",
    "               |-- test.txt\n",
    "               |-- predict.txt\n",
    "\n",
    "Here's a description of the structure:\n",
    "\n",
    "* The ``dataset_root`` directory contains the following:\n",
    "    * ``A``: Contains post-change test images.\n",
    "    * ``B``: Contains pre-change golden reference images.\n",
    "    * ``label``: Contains ground truth segmentation change masks.\n",
    "    * ``list``: Contains .txt files for each dataset split, as described in the  section below. \n",
    "\n",
    "#### List Files\n",
    "\n",
    "\n",
    "Visual ChangeNet-Segmentation dataloader expects the ``label`` directory to contain ``.txt`` files for each of the dataset splits [train, validation, test, inference]. \n",
    "A Visual ChangeNet-Segmentation label file is a ``.txt`` file containing all file names for the particular split.\n",
    "\n",
    "\n",
    "| ``image_names``   | \n",
    "| :---------------- |\n",
    "| ``file_name.png`` |\n",
    "\n",
    "* ``image_names``: The names of images. Image names must be the same for test images and their corresponding reference and mask images. \n",
    "\n",
    "The following sample label file corresponds to the sample directory structure as described.\n",
    "\n",
    "\n",
    "| ``image_names``   | \n",
    "| :---------------- |\n",
    "| ``image1.png``    |\n",
    "| ``image3.png``    |\n",
    "| ``image2.png``    |\n",
    "\n",
    "Note: Each test image (inside directory ``A``) must have a reference image (inside directory ``B``)\n",
    "and a segmentation change map (inside directory ``label``) with the same name for the dataloader to map them correctly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting The MVTech Dataset\n",
    "\n",
    "The MVTech Anomaly dataset provides images of 15 different objects. For each object, the dataset contains non-defective images (golden images) and defective images with a pixel perfect segmentation mask of the defect region. The cells below convert the dataset from any of the 15 objects into the format expected by ChangeNet. \n",
    "\n",
    "* To run this example, go to https://www.mvtec.com/company/research/datasets/mvtec-ad.\n",
    "* Follow the prompts to the download page, right click on any of the object datasets, and copy the link. \n",
    "* Do not copy the link for the whole dataset download because this Notebook only trains on a single object type. \n",
    "* Paste the link in the FIXME below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Download the data\n",
    "import os\n",
    "MVTEC_AD_OBJECT_DOWNLOAD_URL = \"FIXME\"\n",
    "mvtec_object = MVTEC_AD_OBJECT_DOWNLOAD_URL.split(\"/\")[-1].split(\".\")[0]\n",
    "os.environ[\"URL_DATASET\"]=MVTEC_AD_OBJECT_DOWNLOAD_URL\n",
    "os.environ[\"MVTEC_OBJECT\"]=mvtec_object\n",
    "!if [ ! -f $HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz ]; then wget $URL_DATASET -O $HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz; else echo \"image archive already downloaded\"; fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $HOST_DATA_DIR\n",
    "!if [ ! -f $HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz ]; then echo 'Dataset tar file not found, please download.'; else echo 'Found Dataset tar file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!tar xf $HOST_DATA_DIR/$MVTEC_OBJECT.tar.xz -C $HOST_DATA_DIR\n",
    "!chmod -R 755 $HOST_DATA_DIR/$MVTEC_OBJECT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Format the dataset <a class=\"anchor\" id=\"head-2-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install Pillow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random \n",
    "import shutil \n",
    "from PIL import Image\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.environ[\"LOCAL_PROJECT_DIR\"], \"data\", \"changenet\")\n",
    "formatted_dir = f\"formatted_{mvtec_object}_dataset\"\n",
    "\n",
    "DATA_DIR = os.environ[\"HOST_DATA_DIR\"]\n",
    "os.environ[\"FORMATTED_DATA_DIR\"] = formatted_dir\n",
    "\n",
    "#setup dataset folders in expected format \n",
    "formatted_path = os.path.join(DATA_DIR, formatted_dir)\n",
    "a_dir = os.path.join(formatted_path, \"A\")\n",
    "b_dir = os.path.join(formatted_path, \"B\")\n",
    "label_dir = os.path.join(formatted_path, \"label\")\n",
    "list_dir = os.path.join(formatted_path, \"list\")\n",
    "\n",
    "#Create the expected folders\n",
    "os.makedirs(formatted_path, exist_ok=True)\n",
    "os.makedirs(a_dir, exist_ok=True)\n",
    "os.makedirs(b_dir, exist_ok=True)\n",
    "os.makedirs(label_dir, exist_ok=True)\n",
    "os.makedirs(list_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Setup paths to image folders \n",
    "test_dir = os.path.join(DATA_DIR, mvtec_object, \"test\") #contains defect and good images \n",
    "ground_truth_dir = os.path.join(DATA_DIR, mvtec_object, \"ground_truth\") #contains masks for all defect images\n",
    "no_defect_dir = os.path.join(DATA_DIR, mvtec_object, \"train\", \"good\") #contains only good images \n",
    "\n",
    "#Determine defect classes and image size\n",
    "mvtec_object_classes = os.listdir(test_dir)\n",
    "image_name = os.listdir(no_defect_dir)[0] #get an image to determine size\n",
    "mvtec_img_size = Image.open(os.path.join(no_defect_dir,image_name)).size\n",
    "\n",
    "\n",
    "#combine images of different defects into folder 'A'\n",
    "for class_dir in os.listdir(test_dir):\n",
    "    for image in os.listdir(os.path.join(test_dir,class_dir)):\n",
    "        shutil.copy(os.path.join(test_dir,class_dir,image), os.path.join(a_dir, f\"{class_dir}_{image}\"))\n",
    "\n",
    "#combine ground truth images of different defects into folder 'label'\n",
    "for class_dir in os.listdir(test_dir):\n",
    "    for image in os.listdir(os.path.join(test_dir,class_dir)):\n",
    "        if class_dir == \"good\": #make mask for good images\n",
    "            good_mask = Image.new(mode=\"L\", size=mvtec_img_size, color=(0))\n",
    "            good_mask.save(os.path.join(label_dir, f\"{class_dir}_{image}\"))\n",
    "        else: #dataset provides masks for defective images \n",
    "            image_name = os.path.splitext(image)[0]\n",
    "            image_ext = os.path.splitext(image)[1]\n",
    "            src_image = image_name + \"_mask\" + image_ext\n",
    "            shutil.copy(os.path.join(ground_truth_dir,class_dir,src_image), os.path.join(label_dir, f\"{class_dir}_{image}\"))\n",
    "\n",
    "            \n",
    "#Split non defect images to use some for golden samples and some for training \n",
    "golden_split = 0.05 #percentage of good samples from the train set to use as golden images. \n",
    "no_defect_images = os.listdir(no_defect_dir)\n",
    "no_defect_split = int(len(no_defect_images) * golden_split)\n",
    "golden_ref_images = no_defect_images[0:no_defect_split] #taking a % of good images to serve as golden references\n",
    "no_defect_train_images = no_defect_images[no_defect_split:] #adding the rest to the training set \n",
    "\n",
    "#Add good images not used for golden samples into the train set \n",
    "for i, image in enumerate(no_defect_train_images):\n",
    "    dest_image = \"good_2_\" + image\n",
    "    shutil.copy(os.path.join(no_defect_dir, image), os.path.join(a_dir, dest_image))\n",
    "    #also add the mask\n",
    "    good_mask = Image.new(mode=\"L\", size=mvtec_img_size, color=(0))\n",
    "    good_mask.save(os.path.join(label_dir, dest_image))\n",
    "    \n",
    "#Setup the golden reference images \n",
    "for i, image in enumerate(os.listdir(a_dir)):\n",
    "    golden_image = random.sample(golden_ref_images,1)[0] #pair each sample with a golden reference image \n",
    "    shutil.copy(os.path.join(no_defect_dir, golden_image), os.path.join(b_dir, image))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create train, test, validation text files \n",
    "#The datasets are fairly small so we will make test and validation the same\n",
    "\n",
    "train_split = 0.7 #percent to put in train set \n",
    "\n",
    "samples = os.listdir(a_dir)\n",
    "random.shuffle(samples)\n",
    "\n",
    "train = []\n",
    "test = []\n",
    "\n",
    "#Split classes into train and test samples\n",
    "for cls in mvtec_object_classes:\n",
    "    cls_images = [image for image in samples if cls in image]\n",
    "    index_split = int(len(cls_images) * train_split)\n",
    "    train+=cls_images[0:index_split]\n",
    "    test+= cls_images[index_split:]\n",
    "\n",
    "test_labels = \"\\n\".join(test)\n",
    "train_labels = \"\\n\".join(train)\n",
    "\n",
    "#Write out test,val,train lists\n",
    "with open(os.path.join(list_dir, \"test.txt\"), \"w+\") as file:\n",
    "    file.write(test_labels)\n",
    "\n",
    "with open(os.path.join(list_dir, \"val.txt\"), \"w+\") as file:\n",
    "    file.write(test_labels)\n",
    "\n",
    "with open(os.path.join(list_dir, \"train.txt\"), \"w+\") as file:\n",
    "    file.write(train_labels)         "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Verify Dataset Format <a class=\"anchor\" id=\"head-2-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "def verify_dataset(data_dir):\n",
    "\n",
    "    num_images_A = len(os.listdir(a_dir))\n",
    "    num_images_B = len(os.listdir(b_dir))\n",
    "    num_labels = len(os.listdir(label_dir))\n",
    "    num_train_images = sum(1 for _ in open(os.path.join(list_dir, \"train.txt\")))\n",
    "    num_val_images = sum(1 for _ in open(os.path.join(list_dir, \"val.txt\")))\n",
    "    num_test_images = sum(1 for _ in open(os.path.join(list_dir, \"test.txt\")))\n",
    "    print(\"Number of images in the train/val/test set. {}\".format(num_images_A))\n",
    "    print(\"Number of compare images in the train/val/test set. {}\".format(num_images_B))\n",
    "    print(\"Number of labels in the train/val/test set. {}\".format(num_labels))\n",
    "    print(\"Number of train instances. {}\".format(num_train_images))\n",
    "    print(\"Number of val instances. {}\".format(num_val_images))\n",
    "    print(\"Number of test instances. {}\".format(num_test_images))\n",
    "\n",
    "# verify downloaded dataset\n",
    "verify_dataset(data_dir=formatted_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Verify\n",
    "!ls -l $HOST_DATA_DIR/$FORMATTED_DATA_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Download the pretrained model from NGC <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "Download the pretrained model from NGC. The NGC CLI is used to get the data and model. For more details, go to https://ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "import os\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# # Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip\n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_fan_classification_nvimagenet*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model download-version \"nvidia/tao/pretrained_fan_classification_nvimagenet:fan_base_hybrid_nvimagenet\" --dest $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained/pretrained_fan_classification_nvimagenet_vfan_base_hybrid_nvimagenet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Provide training specification <a class=\"anchor\" id=\"head-4\"></a>\n",
    "\n",
    "See the TAO documentation on Visual ChangeNet-Segmentation to determine the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec = \"\"\"\n",
    "encryption_key: tlt_encode\n",
    "task: segment\n",
    "train:\n",
    "  resume_training_checkpoint_path: null\n",
    "  pretrained_model_path: null\n",
    "  segment:\n",
    "    loss: \"ce\"\n",
    "    weights: [0.5, 0.5, 0.5, 0.8, 1.0]\n",
    "  num_epochs: 30\n",
    "  num_nodes: 1\n",
    "  validation_interval: 1\n",
    "  checkpoint_interval: 1\n",
    "  optim:\n",
    "    lr: 0.0002\n",
    "    optim: \"adamw\"\n",
    "    policy: \"linear\" \n",
    "    momentum: 0.9\n",
    "    weight_decay: 0.01\n",
    "results_dir: \"/results\"\n",
    "model:\n",
    "  backbone:\n",
    "    type: \"fan_base_16_p4_hybrid\"\n",
    "    pretrained_backbone_path: /results/pretrained/pretrained_fan_classification_nvimagenet_vfan_base_hybrid_nvimagenet/fan_base_hybrid_nvimagenet.pth\n",
    "dataset:\n",
    "  segment:\n",
    "    dataset: \"CNDataset\"\n",
    "    root_dir: /data/formatted_dataset\n",
    "    data_name: \"custom\"\n",
    "    label_transform: \"norm\"\n",
    "    batch_size: 8\n",
    "    workers: 1\n",
    "    multi_scale_train: True\n",
    "    multi_scale_infer: False\n",
    "    num_classes: 2\n",
    "    img_size: 256\n",
    "    image_folder_name: \"A\"\n",
    "    change_image_folder_name: 'B'\n",
    "    list_folder_name: 'list'\n",
    "    annotation_folder_name: \"label\"\n",
    "    train_split: \"train\"\n",
    "    validation_split: \"val\"\n",
    "    test_split: 'test'\n",
    "    predict_split: 'test'\n",
    "    label_suffix: .png\n",
    "    augmentation: \n",
    "      random_flip:\n",
    "        vflip_probability: 0.5\n",
    "        hflip_probability: 0.5\n",
    "        enable: True\n",
    "      random_rotate:\n",
    "        rotate_probability: 0.5\n",
    "        angle_list: [90, 180, 270]\n",
    "        enable: True\n",
    "      random_color:\n",
    "        brightness: 0.3\n",
    "        contrast: 0.3\n",
    "        saturation: 0.3\n",
    "        hue: 0.3\n",
    "        enable: True\n",
    "      with_scale_random_crop:\n",
    "        enable: True\n",
    "      with_random_crop: True\n",
    "      with_random_blur: True\n",
    "evaluate: \n",
    "  results_dir: \"${results_dir}/evaluate\"\n",
    "  checkpoint: \"${results_dir}/train/changenet.pth\"\n",
    "  trt_engine: \"${gen_trt_engine.trt_engine}\"\n",
    "  batch_size: ${dataset.segment.batch_size}\n",
    "  vis_after_n_batches: 1\n",
    "inference:\n",
    "  results_dir: \"${results_dir}/inference\"\n",
    "  checkpoint: \"${results_dir}/train/changenet.pth\"\n",
    "  trt_engine: \"${gen_trt_engine.trt_engine}\"\n",
    "  batch_size: ${dataset.segment.batch_size}\n",
    "  vis_after_n_batches: 1\n",
    "export:\n",
    "  results_dir: \"${results_dir}/export\"\n",
    "  gpu_id: 0\n",
    "  checkpoint: \"${results_dir}/train/changenet.pth\"\n",
    "  onnx_file: \"${export.results_dir}/changenet.onnx\"\n",
    "  input_width: 256\n",
    "  input_height: 256\n",
    "  batch_size: ${dataset.segment.batch_size}\n",
    "gen_trt_engine:\n",
    "  results_dir: \"${results_dir}/gen_trt_engine\"\n",
    "  gpu_id: 0\n",
    "  onnx_file: \"${export.onnx_file}\"\n",
    "  trt_engine: \"${gen_trt_engine.results_dir}/changenet.trt\"\n",
    "  batch_size: ${dataset.segment.batch_size}\n",
    "  input_channel: 3\n",
    "  input_width: 256\n",
    "  input_height: 256\n",
    "  tensorrt:\n",
    "    data_type: FP32\n",
    "    workspace_size: 1024\n",
    "    min_batch_size: 1\n",
    "    opt_batch_size: 10\n",
    "    max_batch_size: 10\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Write out the spec file \n",
    "with open(os.path.join(os.environ[\"HOST_SPECS_DIR\"], \"experiment.yaml\"), \"w+\") as file:\n",
    "    file.write(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Run TAO training <a class=\"anchor\" id=\"head-5\"></a>\n",
    "* Training will take 20-30 minutes depending on your GPU. \n",
    "\n",
    "* To replicate baseline numbers accurately, the model must be trained for 30 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data/$FORMATTED_DATA_DIR\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results\n",
    "%env NUM_EPOCHS = 30\n",
    "%env BACKBONE_PATH = /results/pretrained/pretrained_fan_classification_nvimagenet_vfan_base_hybrid_nvimagenet/fan_base_hybrid_nvimagenet.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Train Visual ChangeNet-Segmentation model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train a Visual ChangeNet-Segmentation model with a pretrained FAN backbone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train model\")\n",
    "!tao model visual_changenet train \\\n",
    "                  -e $SPECS_DIR/experiment.yaml \\\n",
    "                    train.num_epochs=$NUM_EPOCHS \\\n",
    "                    dataset.segment.root_dir=$DATA_DIR \\\n",
    "                    model.backbone.pretrained_backbone_path=$BACKBONE_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model checkpoints:')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/changenet_model_segment_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/changenet.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/changenet.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Evaluate trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "Evaluate trained model.\n",
    "\n",
    "For Visual ChangeNet-Segmentation model evaluation, the following metrics are preferred:\n",
    "\n",
    "* Overall Accuracy.\n",
    "* mIoU: Intersection over Union (IoU) for all classes.\n",
    "\n",
    "Metrics for Precision, Recall, IoU, and F1 score are provided for each individual class and averaged over all classes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model visual_changenet evaluate \\\n",
    "                   -e $SPECS_DIR/experiment.yaml \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/train/changenet.pth \\\n",
    "                    dataset.segment.root_dir=$DATA_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Visualize Inference <a class=\"anchor\" id=\"head-7\"></a>\n",
    "In this section, run the Visual ChangeNet-Segmentation inference tool to generate inferences with the trained models and save the results under `$RESULTS_DIR`. And visualise a few sample outputs along with their corresponding input images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model visual_changenet inference \\\n",
    "                   -e $SPECS_DIR/experiment.yaml \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/changenet.pth \\\n",
    "                    dataset.segment.root_dir=$DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images(output_path, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"inference\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 8 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Deploy <a class=\"anchor\" id=\"head-8\"></a>\n",
    "Export the model to encrypted ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model visual_changenet export \\\n",
    "                    -e $SPECS_DIR/experiment.yaml \\\n",
    "                        export.checkpoint=$RESULTS_DIR/train/changenet.pth \\\n",
    "                        export.onnx_file=$RESULTS_DIR/export/changenet.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Uncomment this cell to export an onnx file with dynamic batching enabled for\n",
    "# # integration with trtexec and deepstream.\n",
    "# # Export the model with export.batch_size=-1 for dynamic batching.\n",
    "# ! tao model visual_changenet export \\\n",
    "#                     -e $SPECS_DIR/experiment.yaml \\\n",
    "#                     export.checkpoint=$RESULTS_DIR/train/changenet.pth \\\n",
    "#                     export.batch_size=-1 \\\n",
    "#                     export.results_dir=$RESULTS_DIR/export_dynamic \\\n",
    "#                     export.onnx_file=$RESULTS_DIR/export_dynamic/changenet.onnx\n",
    "                            \n",
    "\n",
    "# # Profiling the exported model via trtexec.\n",
    "# ! tao deploy trtexec --onnx=$RESULTS_DIR/export_dynamic/changenet.onnx \\\n",
    "#                      --minShapes=input0:1x3x512x128,input1:1x3x256x256 \\\n",
    "#                      --optShapes=input0:8x3x512x128,input1:8x3x256x256 \\\n",
    "#                      --maxShapes=input0:16x3x512x128,input1:16x3x256x256 \\\n",
    "#                      --fp16 \\\n",
    "#                      --saveEngine=$RESULTS_DIR/export_dynamic/changenet-fp16.engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet gen_trt_engine \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "                gen_trt_engine.onnx_file=$RESULTS_DIR/export/changenet.onnx \\\n",
    "                gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet.trt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet inference \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "                inference.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet.trt \\\n",
    "                inference.results_dir=$RESULTS_DIR/trt_inference \\\n",
    "                dataset.segment.root_dir=$DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet evaluate \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "                evaluate.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet.trt \\\n",
    "                evaluate.results_dir=$RESULTS_DIR/trt_evaluate \\\n",
    "                dataset.segment.root_dir=$DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"trt_inference\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
